"""
Continuity conditions module for locELM.

Handles C^k continuity enforcement across subdomain boundaries.
"""

import tensorflow as tf
import numpy as np
from typing import List, Tuple, Optional
from .domain import Domain, SubdomainInterface
from .networks import MultiSubdomainNetwork


class ContinuityConditions:
    """Enforces C^k continuity conditions across subdomain interfaces."""

    def __init__(self,
                 domain: Domain,
                 networks: MultiSubdomainNetwork,
                 continuity_order: List[int]):
        """
        Initialize continuity conditions.

        Parameters
        ----------
        domain : Domain
            Domain decomposition
        networks : MultiSubdomainNetwork
            Multi-subdomain neural networks
        continuity_order : List[int]
            Continuity order k for each dimension (C^k continuity)
        """
        self.domain = domain
        self.networks = networks
        self.continuity_order = continuity_order
        self.interface_manager = SubdomainInterface(domain)

    def compute_derivatives(self,
                           network_idx: int,
                           points: np.ndarray,
                           order: int,
                           dimension: int) -> np.ndarray:
        """
        Compute derivatives of network output using autodiff.

        Parameters
        ----------
        network_idx : int
            Index of subdomain network
        points : np.ndarray
            Points at which to compute derivatives
        order : int
            Order of derivative
        dimension : int
            Dimension along which to differentiate

        Returns
        -------
        np.ndarray
            Derivative values
        """
        network = self.networks.get_network(network_idx)
        x = tf.constant(points, dtype=tf.float32)

        # Compute derivative using GradientTape
        with tf.GradientTape(persistent=True) as tape:
            tape.watch(x)
            y = network.model(x, training=False)

        if order == 0:
            return y.numpy()

        # First derivative
        dy_dx = tape.gradient(y, x)
        if order == 1:
            return dy_dx[:, dimension:dimension+1].numpy()

        # Second derivative
        with tf.GradientTape() as tape2:
            tape2.watch(x)
            dy_dx_i = tape.gradient(y, x)[:, dimension]

        d2y_dx2 = tape2.gradient(dy_dx_i, x)
        if order == 2:
            return d2y_dx2[:, dimension:dimension+1].numpy()

        raise NotImplementedError("Only up to 2nd order derivatives supported")

    def compute_last_hidden_derivatives(self,
                                        network_idx: int,
                                        points: np.ndarray,
                                        order: int,
                                        dimension: int) -> np.ndarray:
        """
        Compute derivatives of last hidden layer output V_j(x).

        Parameters
        ----------
        network_idx : int
            Index of subdomain network
        points : np.ndarray
            Points at which to compute derivatives
        order : int
            Order of derivative
        dimension : int
            Dimension along which to differentiate

        Returns
        -------
        np.ndarray
            Derivative values of shape (n_points, M)
        """
        network = self.networks.get_network(network_idx)
        x = tf.constant(points, dtype=tf.float32)

        if order == 0:
            return network.compute_last_hidden_output(points)

        # Compute derivatives for each output component
        M = network.hidden_layers[-1]
        n_points = points.shape[0]
        derivatives = np.zeros((n_points, M))

        for j in range(M):
            with tf.GradientTape(persistent=True) as tape:
                tape.watch(x)
                V = network.last_hidden_model(x, training=False)
                V_j = V[:, j]

            if order == 1:
                dV_dx = tape.gradient(V_j, x)
                derivatives[:, j] = dV_dx[:, dimension].numpy()
            elif order == 2:
                with tf.GradientTape() as tape2:
                    tape2.watch(x)
                    dV_dx = tape.gradient(V_j, x)[:, dimension]
                d2V_dx2 = tape2.gradient(dV_dx, x)
                derivatives[:, j] = d2V_dx2[:, dimension].numpy()
            else:
                raise NotImplementedError("Only up to 2nd order derivatives supported")

        return derivatives

    def get_interface_points(self,
                            interface: dict,
                            collocation_points: dict) -> Tuple[np.ndarray, np.ndarray]:
        """
        Get collocation points on an interface from both sides.

        Parameters
        ----------
        interface : dict
            Interface information
        collocation_points : dict
            Dictionary mapping subdomain index to collocation points

        Returns
        -------
        Tuple[np.ndarray, np.ndarray]
            Points from subdomain 1 and subdomain 2 on the interface
        """
        dim = interface['dimension']
        coord = interface['coordinate']

        # Get points from subdomain 1
        points1 = collocation_points[interface['subdomain1']]
        mask1 = np.abs(points1[:, dim] - coord) < 1e-10
        interface_points1 = points1[mask1]

        # Get points from subdomain 2
        points2 = collocation_points[interface['subdomain2']]
        mask2 = np.abs(points2[:, dim] - coord) < 1e-10
        interface_points2 = points2[mask2]

        return interface_points1, interface_points2

    def build_continuity_equations(self,
                                   collocation_points: dict,
                                   output_components: int = 1) -> Tuple[np.ndarray, np.ndarray]:
        """
        Build continuity equations for all interfaces.

        Parameters
        ----------
        collocation_points : dict
            Dictionary mapping subdomain index to collocation points
        output_components : int
            Number of output components

        Returns
        -------
        Tuple[np.ndarray, np.ndarray]
            Coefficient matrix A and right-hand side b for continuity equations
        """
        equations = []

        for interface in self.interface_manager.interfaces:
            dim = interface['dimension']
            k_dim = self.continuity_order[dim]

            points1, points2 = self.get_interface_points(interface, collocation_points)

            # Ensure points match (should be same points on interface)
            n_interface_points = points1.shape[0]

            # Build equations for each order of continuity
            for order in range(k_dim + 1):
                for comp in range(output_components):
                    # Get derivatives from both sides
                    V1 = self.compute_last_hidden_derivatives(
                        interface['subdomain1'], points1, order, dim
                    )
                    V2 = self.compute_last_hidden_derivatives(
                        interface['subdomain2'], points2, order, dim
                    )

                    # Build equation: V1 * w1 - V2 * w2 = 0
                    M = V1.shape[1]
                    n_total_params = self.networks.n_total_params

                    for i in range(n_interface_points):
                        # Create equation row
                        eq_row = np.zeros(n_total_params)

                        # Coefficients for subdomain 1
                        idx1_start = interface['subdomain1'] * M * output_components + comp * M
                        eq_row[idx1_start:idx1_start + M] = V1[i, :]

                        # Coefficients for subdomain 2 (negative)
                        idx2_start = interface['subdomain2'] * M * output_components + comp * M
                        eq_row[idx2_start:idx2_start + M] = -V2[i, :]

                        equations.append(eq_row)

        A = np.array(equations) if equations else np.zeros((0, self.networks.n_total_params))
        b = np.zeros(len(equations))

        return A, b
